import json
import logging
from queue import Queue

from uuid import uuid4
from jaeger_client import Config
import opentracing
import os
from opentracing.propagation import Format


ENCODE = 'utf-8'
DEFUALT_GUID="test_guid"

# 请求上游需要缓存未结束的span
class UpStreamSpanCache():
    span_cache = {}
    tracer = opentracing.tracer

    def __init__(self, tracer):
        self.tracer = tracer

    def put(self, index, span):
        self.span_cache[index] = span

        return True

    def get(self, index):
        return self.span_cache[index] is not None, self.span_cache[index]

    # 请求上游打包span在数据中
    def gen_pynng_req_package(self, flag, span, data):
        span_map = {}
        index = None

        if flag:
            index = str(uuid4())
            self.tracer.inject(span, Format.TEXT_MAP, span_map)
            self.put(index, span)

        package = {
            "flag": flag, 
            "span": span_map,
            "index": index,
            "data": data
        }

        return json.dumps(package).encode(ENCODE)    

    def extract_pynng_resp_package(self, package) -> bytes:
        package = json.loads(package.decode(ENCODE))

        flag = package["flag"]
        if flag:
            success, span = self.get(package["index"])

            if success:
                span.finish()

        return package["data"]   

class DownStreamIndexCache():
    tracer = opentracing.tracer

    def __init__(self, tracer):
        self.tracer = tracer

    # 请求下游从数据包中解析出span 与 data
    def extract_pynng_req_package(self, package):
        package = json.loads(package.decode(ENCODE))

        flag = package["flag"]
        span_ctx = None
        if flag:
            span_ctx = self.tracer.extract(Format.TEXT_MAP, package["span"])

        return flag, package["index"], span_ctx, package["data"]

    # 请求下游生成响应包
    def gen_pynng_resp_package(self, flag, index, data):

        package = {
            "flag": flag, 
            "span": None,
            "index": index,
            "data": data
        }
        
        return json.dumps(package).encode(ENCODE)


def get_guid():
    guid = os.getenv('SUPB_GUID')
    if guid is not None:
        return guid
    else:
        return DEFUALT_GUID

def init_tracer(service):
    logging.getLogger('').handlers = []
    logging.basicConfig(format='%(message)s', level=logging.DEBUG)

    config = Config(
        config={ # usually read from some yaml config
            'sampler': {
                'type': 'const',
                'param': 1,
            },
            'logging': True,
            'reporter_batch_size': 1,
        },
        service_name=service,
    )

    tracer = config.initialize_tracer()
    guid = get_guid()

    logging.debug("use config supb_guid: {}, host: {}, port {}".format(guid, config.local_agent_reporting_host, config.local_agent_reporting_port))

    return tracer
